import argparse

parser = argparse.ArgumentParser()

# train
train_device = ['flyai', '1024gpu', 'PC'][1]
pretrain = True
parser.add_argument("--train-device", default=train_device, type=str, choices=['flyai', '1024gpu', 'PC'],
                    help="training divice")
if train_device == 'flyai':
    parser.add_argument("--train-flyai", default=True, type=bool, help="train in flyai or custom PC")
else:
    parser.add_argument("--train-flyai", default=False, type=bool, help="train in flyai or custom PC")
parser.add_argument("--pretrain", default=pretrain, type=bool, help="load pretrained parameters")
if pretrain:
    parser.add_argument("--pretrain_bald", default=False, type=bool,
                        help="load pretrained parameters on BaldClassification")
parser.add_argument("--unfreeze-epoch", default=9, type=int, help="unfreeze the whole model after some epoch")
parser.add_argument("--multi-model", default=True, type=bool, help="model ensemble")
parser.add_argument("--save-mode", default='ensemble', type=str, choices=['ensemble', 'separate'],
                    help='ensemble: best in the same epoch; separate: best among all batches')

# dataloader
batch_size = [128, 128, 128]
parser.add_argument("-e", "--EPOCHS", default=50, type=int, help="train epochs")
parser.add_argument("-b", "--BATCH", default=batch_size[0], type=int, help="batch size")  # 48, 112
parser.add_argument("--BATCH_train", default=batch_size[0], type=int, help="batch size")
parser.add_argument("--BATCH_val", default=batch_size[1], type=int, help="batch size")
parser.add_argument("--BATCH_test", default=batch_size[2], type=int, help="batch size")
parser.add_argument("--shuffle_train", default=True, type=bool, help="shuffle dataset")
parser.add_argument("--shuffle_val", default=False, type=bool, help="shuffle dataset")
parser.add_argument("--shuffle_test", default=False, type=bool, help="shuffle dataset")

# network
models = ['resnest50', 'densenet121', 'efficientnetb3']
base_lr = [1e-3 / 56, 1.5e-3 / 56, 2.5e-3 / 56]  # lr per epoch
lrs = [base_lr[i] * batch_size[0] for i in range(len(models))]
parser.add_argument("--model1", default=models[0], type=str, choices=models, help="model 1")
parser.add_argument("--model2", default=models[1], type=str, choices=models, help="model 2")
parser.add_argument("--model3", default=models[2], type=str, choices=models, help="model 3")
models = ['resnest50']  # , 'densenet121']
parser.add_argument("--n_models", default=len(models), type=int, help="number of models")
# optimizer & scheduler
parser.add_argument("--lr1", default=lrs[0], type=float, help="learning rate")
parser.add_argument("--eta-min1", default=lrs[0] / 10, type=float, help="min learning rate")
parser.add_argument("--warm-up1", default=5, type=int, help="warm up epoch")
parser.add_argument("--T-01", default=5, type=int, help="half of the first cycle")
parser.add_argument("--T-mult1", default=2, type=int, help="multiply after a cycle")

parser.add_argument("--lr2", default=lrs[1], type=float, help="learning rate of the second net")
parser.add_argument("--eta-min2", default=lrs[1] / 10, type=float, help="min learning rate of the second net")
parser.add_argument("--warm-up2", default=5, type=int, help="warm up epoch of the second net")
parser.add_argument("--T-02", default=5, type=int, help="half of the first cycle of the second net")
parser.add_argument("--T-mult2", default=2, type=int, help="multiply after a cycle of the second net")

parser.add_argument("--lr3", default=lrs[2], type=float, help="learning rate of the third net")
parser.add_argument("--eta-min3", default=lrs[2] / 10, type=float, help="min learning rate of the third net")
parser.add_argument("--warm-up3", default=5, type=int, help="warm up epoch of the second net")
parser.add_argument("--T-03", default=5, type=int, help="half of the first cycle of the second net")
parser.add_argument("--T-mult3", default=2, type=int, help="multiply after a cycle of the second net")

parser.add_argument("--dropout-rate", default=0.5, type=float, help="dropout rate")
parser.add_argument("--freeze-rate", default=0.95, type=float, help="freeze the first several layers")


# loss
# --label smoothing--
parser.add_argument("--do-label-smoothing", default=True, type=bool, help="do label smoothing or not")
parser.add_argument("--label-smoothing", default=0.2, type=float, help="label smoothing parameter")
# --mse loss--
parser.add_argument("--use_mse", default=False, type=bool, help="use MSE loss")
parser.add_argument("--mse_thr", default=0.2, type=float, help="threshold for MSE loss")
# --ArcFace loss--
parser.add_argument("--use_arc", default=False, type=bool, help="use MSE loss")
parser.add_argument("--arc_s", default=25.0, type=float, help="scale factor")
parser.add_argument("--arc_m", default=0.1, type=float, help="margin size")
# --focal loss--
parser.add_argument("--use-focal", default=False, type=bool, help="use focal loss")
parser.add_argument("--f-gamma", default=2.0, type=float, help="gamma parameter")

# optimizer & scheduler
parser.add_argument("--momentum", default=0.9, type=float, help="momentum")
parser.add_argument("--weight-decay", default=3e-5, type=float, help="weight decay")
parser.add_argument("--beta1", default=0.5, type=float, help="Adam: beta1")
parser.add_argument("--beta2", default=0.9, type=float, help="Adam: beta2")
parser.add_argument("--eps", default=1e-8, type=float, help="Adam: eps")
parser.add_argument("--T-max", default=5, type=int, help="half of a cycle")
parser.add_argument("--base_lr", default=1.0, type=float, help="base learning rate")

# dataset
parser.add_argument("--height", default=128, type=int, help="image height")
parser.add_argument("--width", default=128, type=int, help="image width")
parser.add_argument("--train-ratio", default=0.9, type=float, help="ratio of training set")
parser.add_argument("--class-num", default=2, type=int, help="number of class")
parser.add_argument("--shuffle-num", default=10, type=int, help="times to shuffle training data")
parser.add_argument("--train-repeat-time", default=1, type=int, help="times to repeat training data in each epoch")

# --transform--
# ColorJitter
parser.add_argument("--t-brightness", default=0.9, type=float, help="brightness")
parser.add_argument("--t-contrast", default=0.9, type=float, help="contrast")
parser.add_argument("--t-saturation", default=0.99, type=float, help="saturation")
parser.add_argument("--t-hue", default=0.4, type=float, help="hue")
# RandomErasing
parser.add_argument("--t-scale-m", default=0.01, type=float, help="min scale")
parser.add_argument("--t-scale-M", default=0.16, type=float, help="max scale")
parser.add_argument("--t-ratio-m", default=0.4, type=float, help="min ratio")
parser.add_argument("--t-ratio-M", default=2.5, type=float, help="max ratio")
parser.add_argument("--t-value", default='random', type=str, help="value")
# RandomResizedCrop
parser.add_argument("--t-rrc-scale-m", default=0.64, type=float, help="min RandomResizedCrop scale")
parser.add_argument("--t-rrc-scale-M", default=1.0, type=float, help="max RandomResizedCrop scale")

# logging
parser.add_argument("--train-info-interval", default=20, type=int, help="interval to print train info")
parser.add_argument("--num-wrong-imgs", default=5, type=int, help="the number of wrong images to display")
if train_device == 'flyai':
    parser.add_argument("--prefix", default='【', type=str, help="style of print")
    parser.add_argument("--prefix2", default='【', type=str, help="style of print")
    parser.add_argument("--suffix", default='】', type=str, help="style of print")
else:
    parser.add_argument("--prefix", default='\033[1;33;46m', type=str, help="style of print")
    parser.add_argument("--prefix2", default='\033[7;33;46m', type=str, help="style of print")
    parser.add_argument("--suffix", default='\033[0m', type=str, help="style of print")
parser.add_argument("--logging_level", default='DEBUG', type=str,
                    choices=['NOTSET', 'DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'], help='logging level')
parser.add_argument("--print-log-level", default='INFO', type=str,
                    choices=['NOTSET', 'DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'], help='print logging level')

cfg = args = parser.parse_args(args=[])

attrs = [x.dest for x in parser._actions[1:]]  # the same sequence
# [x for x in dir(args) if x[0] != '_']
parser_info = [f"{x}: {eval('args.' + x)}" for x in attrs]
parser_info = '\n'.join(parser_info)
